#!/usr/bin/env python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import abc
import mrget
import shutil
import pickle
import parted
import hashlib
import subprocess


class Util:

    @staticmethod
    def dictForceDelete(the_dict, key):
        if key in the_dict:
            del the_dict[key]

    @staticmethod
    def callRepo(cwd, cmd, argList):
        repoUrl = mrget.target_urls("mirror://repo")[0]
        subprocess.check_call(["repo", cmd, "--repo-url", repoUrl] + argList, cwd=cwd)

    @staticmethod
    def replaceLineInFile(filename, pattern, new_content):
        lineList = []
        bEndWithNewLine = False
        with open(filename, "r") as f:
            buf = f.read()
            lineList = buf.split("\n")
            bEndWithNewLine = buf.endswith("\n")

        for i in range(0, len(lineList)):
            if re.search(pattern, lineList[i]):
                lineList[i] = new_content

        with open(filename, "w") as f:
            f.write("\n".join(lineList))
            if bEndWithNewLine:
                f.write("\n")

    @staticmethod
    def boolTo01(flag):
        return "1" if flag else "0"

    @staticmethod
    def listGetDuplicateItems(theList):
        ret = []
        for i, item in enumerate(theList):
            if item not in ret and item in theList[i+1:]:
                ret.add(item)
        return ret

    @staticmethod
    def listStartswith(theList, subList):
        return len(theList) >= len(subList) and theList[:len(subList)] == subList

    @staticmethod
    def memorySizeStrToInt(sizeStr):
        if sizeStr.endswith("M"):
            return int(sizeStr[:-1]) * 1024 * 1024
        elif sizeStr.endswith("G"):
            return int(sizeStr[:-1]) * 1024 * 1024 * 1024
        else:
            assert False

    @staticmethod
    def diskSizeStrToInt(sizeStr):
        if sizeStr.endswith("M"):
            return int(sizeStr[:-1]) * 1000 * 1000
        elif sizeStr.endswith("G"):
            return int(sizeStr[:-1]) * 1000 * 1000 * 1000
        else:
            assert False

    @staticmethod
    def generateSha256(buf):
        sha256 = hashlib.sha256()
        sha256.update(buf)
        return sha256.hexdigest()

    @staticmethod
    def forceDelete(path):
        if os.path.islink(path):
            os.remove(path)
        elif os.path.isfile(path):
            os.remove(path)
        elif os.path.isdir(path):
            shutil.rmtree(path)
        elif os.path.lexists(path):
            os.remove(path)             # other type of file, such as device node
        else:
            pass                        # path does not exist, do nothing

    @staticmethod
    def forceMakeEmptyDir(path):
        if not os.path.exists(path):
            os.mkdir(path)
        elif os.path.islink(path) or not os.path.isdir(path):
            Util.forceDelete(path)
            os.mkdir(path)
        else:
            for fn in os.listdir(path):
                Util.forceDelete(os.path.join(path, fn))

    @staticmethod
    def removeDirContentExclude(dirPath, excludeList):
        for fn in os.listdir(dirPath):
            if fn not in excludeList:
                Util.forceDelete(os.path.join(dirPath, fn))

    @staticmethod
    def touchFile(filepath):
        assert not os.path.exists(filepath)
        f = open(filepath, 'w')
        f.close()

    @staticmethod
    def moveSoleDirContentOut(dirPath):
        # we process only 1 directory in dirPath
        tlist = os.listdir(dirPath)
        if len(tlist) != 1:
            return False

        # we process only 1 directory in dirPath
        soleDirFn = tlist[0]
        soleDirFullfn = os.path.join(dirPath, soleDirFn)
        if os.path.islink(soleDirFullfn) or not os.path.isdir(soleDirFullfn):
            return False

        # deal with name duplication
        if soleDirFn in os.listdir(soleDirFullfn):
            newFn = "_"
            while len(newFn) <= 10 and newFn in os.list(soleDirFullfn):
                newFn += "_"
            assert len(newFn) <= 10
            newFullfn = os.path.join(dirPath, newFn)
            os.rename(soleDirFullfn, newFullfn)
            soleDirFn = newFn
            soleDirFullfn = newFullfn

        # do work
        for fn in os.listdir(soleDirFullfn):
            os.rename(os.path.join(soleDirFullfn, fn), os.path.join(dirPath, fn))
        os.rmdir(soleDirFullfn)
        return True

    @staticmethod
    def saveObj(filepath, obj):
        with open(filepath, 'wb') as fh:
            pickle.dump(obj, fh)

    @staticmethod
    def loadObj(filepath, klass):
        with open(filepath, "rb") as fh:
            return pickle.load(fh)

    @staticmethod
    def saveEnum(filepath, obj):
        Util.saveObj(filepath, obj)

    @staticmethod
    def loadEnum(filepath, klass):
        return Util.loadObj(filepath)

    @staticmethod
    def readUntil(f, chList, max=-1, bInclusive=False):
        if isinstance(chList, str):
            buf = ""
        elif isinstance(chList, bytes):
            buf = b''
        elif isinstance(chList, list):
            if isinstance(chList[0], str):
                buf = ""
            elif isinstance(chList[0], bytes):
                buf = b''
            else:
                assert False
        else:
            assert False

        while True:
            c = f.read(1)
            if c in chList:
                if bInclusive:
                    buf += c
                break
            buf += c
            if max >= 0 and len(buf) >= max:
                break

        return buf

    @staticmethod
    def isInstanceList(obj, *instances):
        for inst in instances:
            if isinstance(obj, inst):
                return True
        return False

    @staticmethod
    def getMountDeviceForPath(pathname):
        pathname = os.path.abspath(pathname)
        buf = Util.cmdCall("mount")
        for line in buf.split("\n"):
            m = re.search("^(.*) on (.*) type ", line)
            if m is not None and m.group(2) == pathname:
                return m.group(1)
        return None

    @staticmethod
    def mountDiskFile(diskFile, partitionLayout, mountDir, allowRoot=False, partedDisk=None):
        if partedDisk is None:
            partedDisk = parted.newDisk(parted.getDevice(diskFile))
        assert len(partedDisk.partitions) == len(partitionLayout[1])

        mountFileList = []
        try:
            for i, p in enumerate(partedDisk.partitions):
                mountFile = os.path.join(mountDir, "partition%d.loop" % (i))
                Util.touchFile(mountFile)
                cmdList = [
                    "fuseloop",
                    "-O", str(p.geometry.start * partedDisk.device.sectorSize),
                    "-S", str(p.geometry.length * partedDisk.device.sectorSize),
                    diskFile,
                    mountFile,
                ]
                if allowRoot:
                    cmdList += [
                        "-o", "allow_root",
                    ]
                subprocess.run(cmdList, stderr=subprocess.DEVNULL)         # fuseloop prints errors about /etc/mtab
                mountFileList.append(mountFile)
            return mountFileList
        except BaseException:
            for fn in reversed(mountFileList):
                subprocess.run(["umount", mountFile])
                os.unlink(mountFile)
            raise

    @staticmethod
    def umountDiskFile(partitionLayout, mountDir):
        for i in range(0, len(partitionLayout[1])):
            mountFile = os.path.join(mountDir, "partition%d.loop" % (i))
            subprocess.run(["umount", mountFile])
            os.unlink(mountFile)


class TempChdir:

    def __init__(self, dirname):
        self.olddir = os.getcwd()
        os.chdir(dirname)

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        os.chdir(self.olddir)


class ActionRunner:

    class PersistStorage(abc.ABC):

        @abc.abstractmethod
        def initGetCurrentActionInfo(self):
            pass

        @abc.abstractmethod
        def getHistoryActionNames(self):
            pass

        @abc.abstractmethod
        def isFinished(self):
            pass

        @abc.abstractmethod
        def isInUse(self):
            pass

        @abc.abstractmethod
        def use(self):
            pass

        @abc.abstractmethod
        def saveActionStart(self, actionName):
            pass

        @abc.abstractmethod
        def saveActionEnd(self, error=None):
            pass

        @abc.abstractmethod
        def saveFinished(self):
            pass

        @abc.abstractmethod
        def unUse(self):
            pass

    class CustomAction(abc.ABC):

        @abc.abstractmethod
        def get_after(self):
            pass

        @abc.abstractmethod
        def get_before(self):
            pass

    def Action(after=[], before=[], _custom_action_name=None, _custom_action=None):
        def decorator(func):
            def wrapper(self, *kargs, **kwargs):
                curActionIndex = self._getActionIndex(wrapper._action_func_name)
                if self._finished is not None:
                    if self._finished == "":
                        raise self._errClass("build already finished")
                    else:
                        raise self._errClass("build already failed, %s" % (self._finished))
                if curActionIndex != self._lastActionIndex + 1:
                    lastActionName = self._actionList[self._lastActionIndex]._action_func_name[len("action_"):] if self._lastActionIndex >= 0 else "None"
                    raise self._errClass("action must be executed in order (last: %s, current: %s)" % (lastActionName, wrapper._action_func_name[len("action_"):]))
                self._persistStorage.saveActionStart(wrapper._action_func_name[len("action_"):])
                try:
                    if wrapper._action is None:
                        func(self, *kargs, **kwargs)
                    else:
                        func(self, wrapper._action_func_name[len("action_"):], wrapper._action, *kargs, **kwargs)
                except BaseException as e:
                    self._finished = str(e)
                    self._persistStorage.saveActionEnd(self._finished)
                    raise
                else:
                    self._lastActionIndex = curActionIndex
                    self._persistStorage.saveActionEnd()
            wrapper._action_func_name = (func.__name__ if _custom_action_name is None else "action_" + _custom_action_name)
            wrapper._action = _custom_action
            wrapper._after = (after if _custom_action is None else _custom_action.get_after())
            wrapper._before = (before if _custom_action is None else _custom_action.get_before())
            return wrapper
        return decorator

    def __init__(self, persistStorage, actionList, customActionFunc, errClass):
        self._persistStorage = persistStorage
        self._actionList = actionList
        self._endAction = actionList[-1]
        self._customActionFunc = customActionFunc
        self._errClass = errClass

        # check self._actionList
        self._assertActions()

        # check history actions
        historyActionFuncNameList = ["action_" + x for x in self._persistStorage.getHistoryActionNames()]
        actionFuncNameList = [x._action_func_name for x in self._actionList]
        if not Util.listStartswith(actionFuncNameList, historyActionFuncNameList):
            raise self._errClass("invalid history actions")

        # self._lastActionIndex == -1 if no action has been executed
        self._lastActionIndex = len(historyActionFuncNameList) - 1

        # not finished:          self._finished is None
        # successfully finished: self._finished == ""
        # abnormally finished:   self._finished == error-message
        _, err = self._persistStorage.initGetCurrentActionInfo()
        if err is not None:
            self._finished = err
        else:
            self._finished = "" if self._persistStorage.isFinished() else None

        self._persistStorage.use()

    def dispose(self):
        self._persistStorage.unUse()

    def finish(self):
        assert self._finished is None
        assert self._lastActionIndex >= self._actionList.index(self._endAction)
        self._finished = ""
        self._persistStorage.saveFinished()

    def add_custom_action(self, action_name, action, insert_after=None, insert_before=None):
        self.add_custom_actions({action_name: action}, insert_after, insert_before)

    def add_custom_actions(self, action_dict, insert_after=None, insert_before=None):
        for action_name, action in action_dict.items():
            assert re.fullmatch("[0-9a-z_]+", action_name) and "action_" + action_name not in dir(self)
            assert isinstance(action, self.CustomAction)
            assert all([re.fullmatch("[0-9a-z_]+", x) for x in action.get_after()])
            assert all([re.fullmatch("[0-9a-z_]+", x) for x in action.get_before()])

        # convert action object or action name to action index
        if insert_before is not None:
            if insert_before.__class__.__name__ == "method":                        # FIXME
                insert_before = self._actionList.index(insert_before)
            else:
                insert_before = self._getActionIndex("action_" + insert_before)
        if insert_after is not None:
            if insert_after.__class__.__name__ == "method":                         # FIXME
                insert_after = self._actionList.index(insert_after)
            else:
                insert_after = self._getActionIndex("action_" + insert_after)

        # convert to use insert_before only
        if insert_before is not None and insert_after is None:
            pass
        elif insert_before is None and insert_after is not None:
            insert_before = insert_after + 1
        elif insert_before is None and insert_after is None:
            insert_before = len(self._actionList)
        else:
            assert False

        assert self._lastActionIndex < insert_before

        # create new actions and add them to self._actionList
        for action_name, action in action_dict.items():
            func = self.Action(_custom_action_name=action_name, _custom_action=action)(self._customActionFunc)
            func = func.__get__(self)
            exec("self.action_%s = func" % (action_name))
            self._actionList.insert(insert_before, eval("self.action_%s" % (action_name)))
            insert_before += 1

        # do check
        self._assertActions()

    def has_action(self, action_name):
        for action in self._actionList:
            if action._action_func_name == "action_" + action_name:
                return True
        return False

    def add_and_run_custom_action(self, action_name, action):
        self.add_and_run_custom_actions({action_name: action})

    def add_and_run_custom_actions(self, action_dict):
        i = self._lastActionIndex
        for action_name, action in action_dict.items():
            if i == -1:
                self.add_custom_action(action_name, action, insert_before=self._actionList[0])
            else:
                self.add_custom_action(action_name, action, insert_after=self._actionList[i])
            i += 1

        exec("self.action_%s()" % (list(action_dict.keys())[0]))

    def remove_action(self, action_name):
        idx = self._getActionIndex("action_" + action_name)
        assert self._lastActionIndex < idx

        # removes action from self._actionList
        # FIXME: no way to remove action method
        self._actionList.pop(idx)

        # do check
        self._assertActions()

    def get_progress(self):
        if self._finished is None:
            ret = (self._lastActionIndex + 1) * 100 // len(self._actionList)
            return min(ret, 99)
        else:
            return 100

    def _getActionIndex(self, action_func_name):
        for i, action in enumerate(self._actionList):
            if action._action_func_name == action_func_name:
                return i
        assert False

    def _assertActions(self):
        actionFuncNameList = [x._action_func_name for x in self._actionList]
        for i, action in enumerate(self._actionList):
            assert all(["action_" + x not in actionFuncNameList[:i] for x in action._before])
            assert all(["action_" + x not in actionFuncNameList[i+1:] for x in action._after])
